import pdfplumber
from typing import Optional

from ai_quickstart.openai_translator.ai_translator.book import Book, Page, Content, ContentType, TableContent
from ai_quickstart.openai_translator.ai_translator.translator.exceptions import PageOutOfRangeException
from ai_quickstart.openai_translator.ai_translator.utils import LOG


class PDFParser:
    def __init__(self):
        pass

    # 翻译数据准备，将待翻译数据按照 [文本内容、表格内容] -> 页 -> 书的格式逐步汇聚起来，并返回
    # pages：翻译指定页数
    def parse_pdf(self, pdf_file_path: str, pages: Optional[int] = None) -> Book:
        book = Book(pdf_file_path)
        with pdfplumber.open(book.pdf_file_path) as pdf:
            # 指定页数超过原PDF页数时，抛出异常
            if pages is not None and pages > len(pdf.pages):
                raise PageOutOfRangeException(len(pdf.pages), pages)

            # 获取待转换页数，如果指定N，则翻译前N页，如果不指定则翻译全部
            if pages is None:
                pages_to_parse = pdf.pages
            else:
                pages_to_parse = pdf.pages[:pages]

            # 循环获取待翻译页面数据
            for pdf_page in pages_to_parse:
                page = Page()

                # 获取页面文本
                raw_text = pdf_page.extract_text()
                # 获取页面表格列表
                tables = pdf_page.extract_tables()

                for table_data in tables:           # 循环表格
                    for row in table_data:          # 循环表格中的行
                        for cell in row:            # 循环表中中行的列
                            # 替换页面中的文本，去除表格中的对应的text内容，替换不超过1次
                            raw_text = raw_text.replace(cell,'',1)

                if raw_text:
                    # 获取待翻译文本行数据
                    raw_text_lines = raw_text.splitlines()
                    # 获取不为null、空字符串的行
                    cleaned_raw_text_lines = [line.strip() for line in raw_text_lines if line.strip()]
                    # 关联所有行
                    cleaned_raw_text = "\n".join(cleaned_raw_text_lines)

                    # 添加文本内容到页中
                    text_content = Content(content_type=ContentType.TEXT, original=cleaned_raw_text)
                    page.add_content(text_content)
                    LOG.debug(f"[raw_text]\n{cleaned_raw_text}")

                if tables:
                    # 添加表格数据到页中
                    table = TableContent(tables)
                    page.add_content(table)
                    LOG.debug(f"[tables]\n{table}")

                book.add_page(page)
        return book